import time

import torch


def time_it(num_runs=1, *, skipped_runs=1):
    def decorator(func):
        def wrapper(*args, **kwargs):
            result = None
            for _ in range(skipped_runs):
                func(*args, **kwargs)
            total_elapsed_time = 0
            for _ in range(num_runs):
                start_time = time.perf_counter()
                result = func(*args, **kwargs)
                end_time = time.perf_counter()
                elapsed_time = end_time - start_time
                total_elapsed_time += elapsed_time
            average_elapsed_time = total_elapsed_time / num_runs
            print(
                f"函数 {func.__name__} 跳过 {skipped_runs} 次后执行 {num_runs} 次的平均耗时：{average_elapsed_time * 1000000}us")
            return result

        return wrapper

    return decorator


@time_it(num_runs=3)
@torch.compile(dynamic=True)
def llama2_patten1(x1, x2):
    x1 = torch.mul(x1, torch.sigmoid(x1))
    return torch.mul(x2, x1)


x1 = torch.ones(22016, 512)
x2 = torch.ones(22016, 512)
y = llama2_patten1(x1, x2)
print(y)
